{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": [
     "notebook-header"
    ]
   },
   "source": [
    "[![GitHub issues by-label](https://img.shields.io/github/issues-raw/pfebrer/sisl/WavefunctionPlot?style=for-the-badge)](https://github.com/pfebrer/sisl/labels/WavefunctionPlot)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "WavefunctionPlot\n",
    "=========\n",
    "\n",
    "The `WavefunctionPlot` class will help you very easily generate and display wavefunctions from a `Hamiltonian` or any other source. If you already have your wavefunction in a grid, you can use `GridPlot`.  \n",
    "\n",
    "<div class=\"alert alert-info\">\n",
    "    \n",
    "Note\n",
    "    \n",
    "`WavefunctionPlot` is just an extension of `GridPlot`, so everything in [the GridPlot notebook](./GridPlot.ipynb) applies and this notebook **will only display the additional features**.\n",
    "\n",
    "`WavefunctionPlot` changes the defaults of the `axes` and `plot_geom` inputs so that by default the grid is shown in 3D and displaying the geometry.\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sisl"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generating wavefunctions from a hamiltonian\n",
    "\n",
    "We will create a toy graphene tight binding hamiltonian, but you could have read the `Hamiltonian` from any source. Note that your hamiltonian **needs to contain the corresponding geometry with the right orbitals**, otherwise we have no idea what's the shape of the wavefunction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "r = np.linspace(0, 3.5, 50)\n",
    "f = np.exp(-r)\n",
    "\n",
    "orb = sisl.AtomicOrbital(\"2pzZ\", (r, f))\n",
    "geom = sisl.geom.graphene(orthogonal=True, atoms=sisl.Atom(6, orb))\n",
    "geom = geom.move([0, 0, 5])\n",
    "H = sisl.Hamiltonian(geom)\n",
    "H.construct(\n",
    "    [(0.1, 1.44), (0, -2.7)],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we have our hamiltonian, plotting a wavefunction is as simple as: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "H.plot.wavefunction()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is a very delocalized state, so its representation in 3D is not very interesting"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Selecting the wavefunction\n",
    "\n",
    "By default, `WavefunctionPlot` gives you the first wavefunction at the gamma point. You can control this behavior by tuning the `i` and `k` settings.\n",
    "\n",
    "For example, to get the second wavefunction at the gamma point:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot = H.plot.wavefunction(i=2, k=(0, 0, 0))\n",
    "plot"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can also select the spin with the `spin` setting (if you have, of course, a spin polarized `Hamiltonian`).\n",
    "\n",
    "<div class=\"alert alert-info\">\n",
    "\n",
    "Note\n",
    "    \n",
    "If you update the **number of the wavefunction, the eigenstates are already calculated**, so there's no need to recalculate them. However, changing the **k point** or the **spin component** will trigger a **recalculation of the eigenstates**.\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Grid precision\n",
    "\n",
    "The wavefunction is projected in a grid, and how fine that grid is will determine the resolution. You can control this with the `grid_prec` setting, which accepts the **grid precision in Angstrom**. Let's check the difference in 2D, where it will be best appreciated:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot.update_inputs(axes=\"xy\", transforms=[\"square\"])  # by default grid_prec is 0.2 Ang"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot.update_inputs(grid_prec=0.05)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Much better, isn't it? Notice how it didn't look that bad in 3d, because the grid is smooth, so it's values are nicely interpolated. You can also appreciate this by setting `smooth` to `True` in 2D, which does an \"OK job\" at guessing the values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot.update_inputs(grid_prec=0.2, smooth=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-warning\">\n",
    "    \n",
    "Warning\n",
    "    \n",
    "Keep in mind that a finer grid will **occupy more memory and take more time to generate and render**, and sometimes it might be unnecessary to make your grid very fine, specially if it's smooth.\n",
    "    \n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## GridPlot settings\n",
    "\n",
    "As stated at the beggining of this notebook, you have all the power of `GridPlot` available to you. Therefore you can, for example, display supercells of the resulting wavefunctions (please don't tile the hamiltonian! :))."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "nbsphinx-thumbnail"
    ]
   },
   "outputs": [],
   "source": [
    "plot.update_inputs(\n",
    "    axes=\"xyz\",\n",
    "    nsc=[2, 2, 1],\n",
    "    grid_prec=0.1,\n",
    "    transforms=[],\n",
    "    isos=[\n",
    "        {\"val\": -0.07, \"opacity\": 1, \"color\": \"salmon\"},\n",
    "        {\"val\": 0.07, \"opacity\": 0.7, \"color\": \"blue\"},\n",
    "    ],\n",
    "    geom_kwargs={\"atoms_style\": dict(color=[\"orange\", \"red\", \"green\", \"pink\"])},\n",
    ").show(\"png\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": [
     "notebook-end"
    ]
   },
   "source": [
    "We hope you enjoyed what you learned!"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
